import atexit
import multiprocessing as mp
import time


# 1. 将Config类定义在全局作用域（子进程可访问）
class Config:
    tensor_parallel_size = 3  # 启动2个工作进程（1和2）


def ModelRunner(config, index, event):
    """子进程执行的任务：等待事件触发后运行"""
    print(f"工作进程 {index} 启动，等待同步事件...")
    event.wait()  # 阻塞，直到主进程触发事件
    print(f"工作进程 {index} 开始执行任务（tensor_parallel_size={config.tensor_parallel_size}）")
    print(f"工作进程 {index} 任务完成")


class ProcessManager:
    def __init__(self, config):
        self.ps = []  # 子进程列表
        self.events = []  # 同步事件列表
        ctx = mp.get_context("spawn")  # 跨平台安全的进程创建方式

        # 启动工作进程（从1到tensor_parallel_size-1）
        for i in range(1, config.tensor_parallel_size):
            event = ctx.Event()  # 创建同步事件
            # 传递config实例给子进程（子进程可访问Config类）
            process = ctx.Process(target=ModelRunner, args=(config, i, event))
            process.start()
            self.ps.append(process)
            self.events.append(event)

    def trigger_events(self):
        """触发所有事件，让子进程开始执行"""
        print("主进程：触发所有同步事件")
        for event in self.events:
            event.set()

    def join_processes(self):
        """等待所有子进程结束"""
        for process in self.ps:
            process.join()

def exit():
    print(f"exit")

if __name__ == "__main__":
    # 2. 在主模块中创建Config实例（子进程可识别）
    config = Config()
    atexit.register(exit)
    manager = ProcessManager(config)

    print("主进程：等待3秒后触发事件...")
    time.sleep(3)  # 模拟主进程准备工作

    manager.trigger_events()  # 通知子进程开始执行
    manager.join_processes()  # 等待所有子进程完成
    print("主进程：所有子进程已完成")
